from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
import os
import base64


class AES256Cipher:
    """
    AES-256 加密工具类 (使用 cryptography 库实现)

    功能特点：
    1. 从自定义字符串生成密钥
    2. 使用 CBC 模式
    3. 自动处理 PKCS7 填充
    4. 每次加密生成随机 IV
    5. 加密结果包含 IV，便于解密

    使用示例：
        cipher = AES256Cipher("我的秘密密码")
        encrypted = cipher.encrypt("要加密的消息")
        decrypted = cipher.decrypt(encrypted)
    """

    def __init__(self, password: str, salt: bytes = None, iterations: int = 100000):
        """
        初始化加密工具

        :param password: 用于生成密钥的自定义字符串
        :param salt: 可选的盐值，如果不提供则随机生成
        :param iterations: PBKDF2 密钥派生迭代次数
        """
        if salt is None:
            self.salt = os.urandom(16)
        else:
            self.salt = salt

        # 使用 PBKDF2 从密码派生密钥
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,  # AES-256 需要 32 字节密钥
            salt=self.salt,
            iterations=iterations,
            backend=default_backend()
        )
        self.key = kdf.derive(password.encode('utf-8'))

    def encrypt(self, plaintext: str) -> str:
        """
        加密文本

        :param plaintext: 要加密的文本
        :return: base64 编码的加密结果（包含 IV）
        """
        # 生成随机 IV
        iv = os.urandom(16)

        # 创建加密器
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.CBC(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()

        # 应用 PKCS7 填充
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(plaintext.encode('utf-8')) + padder.finalize()

        # 加密
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()

        # 组合 IV 和密文
        combined = iv + ciphertext

        # 返回 base64 编码结果
        return base64.b64encode(combined).decode('utf-8')

    def decrypt(self, encrypted_data: str) -> str:
        """
        解密文本

        :param encrypted_data: base64 编码的加密数据（包含 IV）
        :return: 解密后的原始文本
        """
        # 解码 base64
        combined = base64.b64decode(encrypted_data)

        # 提取 IV 和密文
        iv = combined[:16]
        ciphertext = combined[16:]

        # 创建解密器
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.CBC(iv),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()

        # 解密
        padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()

        # 去除 PKCS7 填充
        unpadder = padding.PKCS7(128).unpadder()
        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()

        return plaintext.decode('utf-8')

    def get_salt(self) -> str:
        """
        获取使用的盐值（base64 编码）
        """
        return base64.b64encode(self.salt).decode('utf-8')


# 使用示例
# if __name__ == "__main__":
#     # 示例 1：使用随机盐
#     password = "这是我的秘密密码"
#     cipher1 = AES256Cipher(password)
#     print("使用的盐值:", cipher1.get_salt())
#
#     message = "这是一条需要加密的重要消息"
#     print("原始消息:", message)
#
#     encrypted = cipher1.encrypt(message)
#     print("加密结果:", encrypted)
#
#     decrypted = cipher1.decrypt(encrypted)
#     print("解密结果:", decrypted)
#
#     # 示例 2：使用固定盐（便于测试，实际应用中不推荐）
#     salt = b'fixed_salt_12345'  # 实际应用中应该使用随机盐
#     cipher2 = AES256Cipher(password, salt=salt)
#
#     print("\n使用固定盐的示例:")
#     encrypted2 = cipher2.encrypt(message)
#     print("加密结果:", encrypted2)
#     decrypted2 = cipher2.decrypt(encrypted2)
#     print("解密结果:", decrypted2)